from flask import Blueprint, request, jsonify
import os
import subprocess
import requests
from datetime import datetime
from werkzeug.utils import secure_filename
from utils.logger import get_logger
from config import Config
from database import TaskDAO
import shutil

logger = get_logger(__name__)

def create_training_routes(task_dao: TaskDAO):
    """创建训练相关路由"""

    training_bp = Blueprint('training', __name__)

    # 配置
    TRAINING_SERVER_URL = "http://192.168.0.152:2000"
    TTS_SERVER_URL = "http://192.168.0.152:18180"

    @training_bp.route('/start-training', methods=['POST'])
    def start_training():
        """开始训练任务"""
        try:
            data = request.get_json()
            task_id = data.get('taskId')
            video_path = data.get('videoPath')
            primary_category_id = data.get('primaryCategoryId')
            secondary_category_id = data.get('secondaryCategoryId')
            tertiary_category_id = data.get('tertiaryCategoryId', 0)
            model_name = data.get('modelName', f'Model_{datetime.now().strftime("%Y%m%d%H%M%S")}')

            # 验证参数
            if not all([task_id, video_path, primary_category_id, secondary_category_id]):
                return jsonify({'success': False, 'message': '缺少必要参数'}), 400

            # 检查视频文件是否存在
            if not os.path.exists(video_path):
                return jsonify({'success': False, 'message': '视频文件不存在'}), 404

            # 1. 生成文件名
            timestamp = datetime.now().strftime('%Y%m%d%H%M%S%f')[:-3]  # 毫秒级时间戳
            video_ext = os.path.splitext(video_path)[1]
            model_file_name = f"{timestamp}{video_ext}"
            audio_file_name = f"{timestamp}.wav"

            # 2. 提取音频到临时目录
            logger.info(f"开始提取音频: {video_path}")
            temp_audio_path = os.path.join(Config.TEMP_FOLDER, audio_file_name)

            if not extract_audio_from_video(video_path, temp_audio_path):
                return jsonify({'success': False, 'message': '音频提取失败'}), 500

            # 3. 复制视频到临时目录（用于上传）
            temp_video_path = os.path.join(Config.TEMP_FOLDER, model_file_name)
            shutil.copy2(video_path, temp_video_path)

            # 4. 上传视频到训练服务器
            logger.info("上传视频到训练服务器")
            video_upload_result = upload_model_video(temp_video_path, model_file_name, TRAINING_SERVER_URL)
            if not video_upload_result:
                cleanup_temp_files([temp_video_path, temp_audio_path])
                return jsonify({'success': False, 'message': '视频上传失败'}), 500

            # 5. 上传音频到训练服务器
            logger.info("上传音频到训练服务器")
            audio_upload_result = upload_audio(temp_audio_path, audio_file_name, TRAINING_SERVER_URL)
            if not audio_upload_result:
                cleanup_temp_files([temp_video_path, temp_audio_path])
                return jsonify({'success': False, 'message': '音频上传失败'}), 500

            # 6. 设置训练路径（使用相对路径）
            train_path = f"origin_audio/{audio_file_name}"

            # 7. 调用训练接口
            logger.info(f"开始训练语音模型，训练路径: {train_path}")
            voice_id = train_voice(train_path, 'zh', TRAINING_SERVER_URL, TTS_SERVER_URL)
            if not voice_id:
                cleanup_temp_files([temp_video_path, temp_audio_path])
                return jsonify({'success': False, 'message': '语音训练失败'}), 500

            logger.info(f"语音模型训练完成，voiceId: {voice_id}")

            # 8. 设置相对路径（用于数据库存储）
            relative_model_path = model_file_name
            relative_audio_path = f"origin_audio\\{audio_file_name}"  # 使用反斜杠以匹配原始代码

            # 9. 将模型信息同步到服务器
            logger.info("同步模型信息到服务器")
            sync_result = sync_model_to_server({
                'modelName': model_name,
                'videoPath': relative_model_path,
                'audioPath': relative_audio_path,
                'imagePath': video_upload_result.get('thumbnailUrl', ''),
                'voiceId': voice_id,
                'primaryCategoryId': primary_category_id,
                'secondaryCategoryId': secondary_category_id,
                'tertiaryCategoryId': tertiary_category_id,
                'lang': 'zh'
            }, TRAINING_SERVER_URL)

            if not sync_result:
                cleanup_temp_files([temp_video_path, temp_audio_path])
                return jsonify({'success': False, 'message': '模型同步失败'}), 500

            # 10. 更新任务状态为已训练
            task_dao.update_task_trained_status(task_id, True)

            # 11. 清理临时文件
            cleanup_temp_files([temp_video_path, temp_audio_path])

            return jsonify({
                'success': True,
                'message': '训练任务已成功提交',
                'voiceId': voice_id
            })

        except Exception as e:
            logger.error(f"训练任务失败: {e}")
            return jsonify({'success': False, 'message': str(e)}), 500

    def extract_audio_from_video(video_path, audio_path):
        """从视频中提取音频"""
        try:
            # 确保输出目录存在
            os.makedirs(os.path.dirname(audio_path), exist_ok=True)

            # 使用ffmpeg提取音频
            cmd = [
                'ffmpeg',
                '-i', video_path,
                '-vn',  # 不包含视频
                '-acodec', 'pcm_s16le',  # WAV格式
                '-ar', '44100',  # 采样率
                '-ac', '2',  # 立体声
                '-y',  # 覆盖输出文件
                audio_path
            ]

            # logger.info(f"执行FFmpeg命令: {' '.join(cmd)}")
            result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='replace')

            if result.returncode != 0:
                logger.error(f"FFmpeg音频提取失败: {result.stderr}")
                return False

            logger.info(f"音频提取成功: {audio_path}")
            return True

        except Exception as e:
            logger.error(f"音频提取异常: {e}")
            return False

    def upload_model_video(video_path, filename, server_url):
        """上传模型视频"""
        try:
            url = f"{server_url}/apis/upload/upload/model"

            with open(video_path, 'rb') as f:
                files = {
                    'file': (filename, f, 'video/mp4')
                }
                # 添加filename参数
                data = {'filename': filename}

                response = requests.post(url, files=files, data=data)

            if response.status_code == 200:
                result = response.json()
                if result.get('code') == 0:
                    return {
                        'videoPath': result['data']['videoPath'],
                        'thumbnailUrl': result['data'].get('thumbnailUrl', '')
                    }

            logger.error(f"视频上传失败: {response.text}")
            return None

        except Exception as e:
            logger.error(f"视频上传异常: {e}")
            return None

    def upload_audio(audio_path, filename, server_url):
        """上传音频文件"""
        try:
            url = f"{server_url}/apis/upload/upload/audio"

            with open(audio_path, 'rb') as f:
                files = {
                    'file': (filename, f, 'audio/wav')
                }
                # 添加filename参数
                data = {'filename': filename}

                response = requests.post(url, files=files, data=data)

            if response.status_code == 200:
                result = response.json()
                if result.get('code') == 0:
                    return {
                        'path': result['data']
                    }

            logger.error(f"音频上传失败: {response.text}")
            return None

        except Exception as e:
            logger.error(f"音频上传异常: {e}")
            return None

    def train_voice(audio_path, lang, server_url, tts_url):
        """训练语音模型"""
        try:
            # 将路径中的反斜杠替换为正斜杠
            audio_path = audio_path.replace('\\', '/')

            # 第一步：调用TTS服务进行预处理
            logger.info(f"调用TTS服务预处理音频: {audio_path}")

            preprocess_data = {
                'format': audio_path.split('.')[-1],  # 获取文件格式
                'reference_audio': audio_path,
                'lang': lang
            }

            response = requests.post(
                f"{tts_url}/v1/preprocess_and_tran",
                json=preprocess_data
            )

            if response.status_code != 200:
                logger.error(f"TTS预处理失败，状态码: {response.status_code}")
                return None

            result = response.json()
            # logger.info(f"TTS预处理响应: {result}")

            if result.get('code') != 0:
                logger.error(f"TTS预处理失败: {result}")
                return None

            # 获取预处理结果
            asr_format_audio_url = result.get('asr_format_audio_url')
            reference_audio_text = result.get('reference_audio_text')

            # 第二步：添加语音模型到主服务器
            logger.info("添加语音模型到服务器")

            voice_data = {
                'origin_audio_path': audio_path,
                'lang': lang,
                'asr_format_audio_url': asr_format_audio_url,
                'reference_audio_text': reference_audio_text
            }

            add_response = requests.post(
                f"{server_url}/apis/model/api/voice/add",
                json=voice_data
            )

            if add_response.status_code != 200:
                logger.error(f"添加语音模型失败，状态码: {add_response.status_code}")
                return None

            add_result = add_response.json()
            # logger.info(f"添加语音模型响应: {add_result}")

            # 返回语音模型ID
            if add_result.get('data') and add_result['data'].get('id'):
                return add_result['data']['id']

            logger.error("未能获取语音模型ID")
            return None

        except Exception as e:
            logger.error(f"语音训练异常: {e}")
            return None

    def sync_model_to_server(model_data, server_url):
        """同步模型到服务器"""
        try:
            url = f"{server_url}/apis/model/api/model/add"

            # logger.info(f"同步模型数据: {model_data}")
            response = requests.post(url, json=model_data)

            if response.status_code == 200:
                result = response.json()
                # logger.info(f"模型同步响应: {result}")
                return result.get('code') == 200

            logger.error(f"模型同步失败: {response.text}")
            return False

        except Exception as e:
            logger.error(f"模型同步异常: {e}")
            return False

    def cleanup_temp_files(file_paths):
        """清理临时文件"""
        for file_path in file_paths:
            try:
                if os.path.exists(file_path):
                    os.remove(file_path)
                    logger.info(f"已清理临时文件: {file_path}")
            except Exception as e:
                logger.error(f"清理文件失败 {file_path}: {e}")

    # 分类相关API代理
    @training_bp.route('/category/tree', methods=['GET'])
    def get_category_tree():
        """获取分类树"""
        try:
            response = requests.get(f"{TRAINING_SERVER_URL}/apis/category/api/category/tree")
            return jsonify(response.json())
        except Exception as e:
            logger.error(f"获取分类树失败: {e}")
            return jsonify({'code': 500, 'message': '获取分类树失败'})

    @training_bp.route('/category/sub/<int:parent_id>', methods=['GET'])
    def get_sub_categories(parent_id):
        """获取子分类"""
        try:
            response = requests.get(f"{TRAINING_SERVER_URL}/apis/category/api/category/sub/{parent_id}")
            return jsonify(response.json())
        except Exception as e:
            logger.error(f"获取子分类失败: {e}")
            return jsonify({'code': 500, 'message': '获取子分类失败'})

    @training_bp.route('/category/add', methods=['POST'])
    def add_category():
        """添加分类"""
        try:
            data = request.get_json()
            response = requests.post(
                f"{TRAINING_SERVER_URL}/apis/category/api/category/add",
                json=data
            )
            return jsonify(response.json())
        except Exception as e:
            logger.error(f"添加分类失败: {e}")
            return jsonify({'code': 500, 'message': '添加分类失败'})

    @training_bp.route('/category/delete/<int:category_id>', methods=['DELETE'])
    def delete_category(category_id):
        """删除分类"""
        try:
            response = requests.delete(
                f"{TRAINING_SERVER_URL}/apis/category/api/category/delete/{category_id}"
            )
            return jsonify(response.json())
        except Exception as e:
            logger.error(f"删除分类失败: {e}")
            return jsonify({'code': 500, 'message': '删除分类失败'})

    return training_bp